#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
【节点名称】：
    CameraCalibrationNode
【依赖项安装】：
    pip install spirems
    git clone -b spirecv2-dds https://gitee.com/amovlab/SpireCV.git
    cd SpireCV && pip install -e .
【订阅类型】：
    sensor_msgs::CompressedImage （输入图像）
【发布类型】：
    暂无数据发布
  
【构造参数说明】：
    parameter_file (str): 全局参数文件
    sms_shutdown (bool): 是否接收全局关闭信号,如果需要长期后台执行,建议设置为False
    specified_input_topic (str): 指定输入的话题地址
    realtime_det (bool): 是否是实时,设置为True可以降低延迟,但可能会产生丢帧
【节点参数】：
    square_size : 图像标定板的方块真实宽度, 默认 10 (mm)
    output_file : 标定结果文件名称，默认 calib.yaml
【备注】：
    无
"""


import numpy as np
import cv2 as cv
# local modules
from common import splitfn
# built-in modules
import yaml 
import threading
from spirems import Subscriber, sms2cvimg, Logger, BaseNode, get_extra_args
from queue import Queue
from typing import Union
import json
import sys
import argparse
import time
import os


obj_points = []
img_points = []



def save_camera_params_to_json(filename, imageSize, flags, cameraMatrix, distCoeffs):
    # 获取当前时间并格式化为字符串
    calibration_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    cameraMatrix_f = cameraMatrix.flatten()
    distCoeffs_f = distCoeffs.flatten()

    # 这里的相机外参默认全为0
    R = np.identity(3)
    t = np.zeros((3, 1))
    extrinsicMatrix = np.hstack((R, t))
    projectionMatrix = cameraMatrix @ extrinsicMatrix
    projectionMatrix_f = projectionMatrix.flatten()

    # 准备一个字典来存储相机参数
    camera_params ={}
    camera_params["/CameraNode/calibration_time"] = calibration_time
    camera_params["/CameraNode/image_width"] = imageSize[0]
    camera_params["/CameraNode/image_height"] = imageSize[1]
    camera_params["/CameraNode/flags"] = flags
    camera_params["/CameraNode/flags_description"] =  _flags_to_description(flags)
    camera_params["/CameraNode/camera_matrix"] =  cameraMatrix_f.tolist()
    camera_params["/CameraNode/distortion_coefficients"] = distCoeffs_f.tolist()
    camera_params["/CameraNode/projection"] = projectionMatrix_f.tolist()

    # 写入 JSON 文件
    with open(filename, 'w') as file:
        json.dump(camera_params, file, indent=4)


def _flags_to_description(flags):
    # 将标志位转换为描述字符串
    flag_strings = []
    if flags & cv.CALIB_USE_INTRINSIC_GUESS:
        flag_strings.append("use_intrinsic_guess")
    if flags & cv.CALIB_FIX_ASPECT_RATIO:
        flag_strings.append("fix_aspect_ratio")
    if flags & cv.CALIB_FIX_PRINCIPAL_POINT:
        flag_strings.append("fix_principal_point")
    if flags & cv.CALIB_ZERO_TANGENT_DIST:
        flag_strings.append("zero_tangent_dist")
    return ", ".join(flag_strings)


def calibrate(img, square_size, imw, imh, patternw, patternh, output_file):
    t1 = time.time()
    square_size = float(square_size)
    pattern_size = (patternw, patternh)
    pattern_points = np.zeros((np.prod(pattern_size), 3), np.float32)
    pattern_points[:, :2] = np.indices(pattern_size).T.reshape(-1, 2)
    pattern_points *= square_size
    output_file = str(output_file)

    img = cv.resize(img, (imw, imh))
    frame = img

    cv.putText(frame, "Press 'c' to add current frame. 'ESC' to finish and calibrate", (10, 20), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2, lineType=cv.LINE_AA)
    found, corners = cv.findChessboardCorners(img, pattern_size)
    if found:
        img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
        term = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_COUNT, 30, 0.1)
        cv.cornerSubPix(img, corners, (5, 5), (-1, -1), term)
        cv.drawChessboardCorners(frame, pattern_size, corners, found)

    key = cv.waitKey(1) & 0xFF

    if key == ord('c'):
        # 将当前帧添加到截图列表中
        if corners is not None:
            print("Frame captured.")
            img_points.append(corners.reshape(-1, 2))
            obj_points.append(pattern_points)

    # 检查是否按下退出键（例如 'ESC'）
    elif key == 27:
        if len(img_points) < 4:
            print("Not enough captures for calibration")
        else:
            # calculate camera distortion
            rms, camera_matrix, dist_coefs, _rvecs, _tvecs = cv.calibrateCamera(obj_points, img_points, (imw, imh), None, None)

            flags = cv.CALIB_FIX_ASPECT_RATIO | cv.CALIB_USE_INTRINSIC_GUESS
            save_camera_params_to_json(output_file, (imw, imh), flags, camera_matrix, dist_coefs)

            print("\nRMS:", rms)
            print("camera matrix:\n", camera_matrix)
            print("distortion coefficients: ", dist_coefs.ravel())
            print("Calibration saved to : ", output_file)

            print('Done!')
        frame = None
    print("FPS: {:.2f}".format(1.0 / (time.time() - t1)))
    return frame


class CameraCalibrationNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self, 
            self.__class__.__name__, 
            job_name, 
            ip=ip, 
            port=port, 
            param_dict_or_file=param_dict_or_file, 
            sms_shutdown=sms_shutdown,
            **kwargs
        )

        self.image_topic = self.get_param('image_topic', '/' + job_name + '/sensor/image_raw')
        self.square_size = self.get_param('square_size', 10.0)
        self.pattern_size = self.get_param('pattern_size', [9, 6])
        self.output_file = self.get_param('output_file', 'calib.json')
        self.realtime_det = self.get_param('realtime_det', True)
        self.params_help()

        self._image_reader = Subscriber(
            self.image_topic, 'std_msgs::Null', self.image_callback,
            ip=ip, port=port
        )

        self.job_queue = Queue()
        self.queue_pool.append(self.job_queue)
        self.logger = Logger(self.__class__.__name__)

        self.logger.info("image_topic: " + self.image_topic)
        self.logger.info("square_size: " + str(self.square_size))
        self.logger.info("pattern_size: " + str(self.pattern_size))
        self.logger.info("output_file: " + self.output_file)
        self.logger.info("realtime_det: " + str(self.realtime_det))

        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()
        self.logger.quit()

    def image_callback(self, msg):
        if self.realtime_det:
            if not self.job_queue.empty():
                self.job_queue.queue.clear()
        img = sms2cvimg(msg)
        self.job_queue.put({'msg': msg, 'img': img})

    def run(self):
        while self.is_running():
            msg_img = self.job_queue.get(block=True)
            if msg_img is None:
                break

            img, msg = msg_img['img'], msg_img['msg']
            frame = calibrate(img, self.square_size, img.shape[1], img.shape[0], self.pattern_size[0], self.pattern_size[1], self.output_file)

            if frame is None:
                break
            else:
                cv.imshow('out', frame)

        cv.destroyAllWindows()
        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("config:", args.config)
    print("job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = CameraCalibrationNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.join() 
